# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
from typing import List
from tortoise import fields, Tortoise
from kernels.model import DbModel


class PgKnowledgeModel(DbModel):
    uuid = fields.UUIDField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="拥有者ID")
    last_uid = fields.IntField(null=False, default=0, description="操作者ID")
    know_id = fields.IntField(null=False, default=0, description="知识库ID")
    archive_id = fields.IntField(null=False, default=0, description="关联文档ID")
    vector_model = fields.CharField(null=False, max_length=200, default="", description="向量模型")
    vector_alias = fields.CharField(null=False, max_length=200, default="", description="向量别名")
    metadata = fields.TextField(default="", description="元信息json")
    signed = fields.CharField(null=False, max_length=32, default="", description="唯一签名")
    question = fields.TextField(default="", description="问题内容")
    answer = fields.TextField(default="", description="补充内容")
    phrases = fields.TextField(description="中文分词")
    embedding = fields.TextField(description="向量内容")
    dimension = fields.IntField(null=False, default=0, description="向量维度")
    chunk_index = fields.IntField(null=False, default=0, description="分片下标")
    use_points = fields.IntField(null=False, default=0, description="消耗积分")
    use_tokens = fields.IntField(null=False, default=0, description="消耗积分")
    error = fields.TextField(null=False, default="", description="错误信息")
    status = fields.SmallIntField(null=False, default=0, description="学习状态: [1=排队中, 2=学习中, 3=成功, 4=失败]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "向量知识表"
        table = DbModel.table_prefix("knowledge")

    @classmethod
    async def update_ts_vectory(cls, uuids: List[str]):
        """ 更新分词信息 """
        if uuids:
            ids: str = ",".join(f"'{x}'" for x in uuids)
            SQL = cls.filter(uuid__in=[""]).update(phrases="").sql()
            SQL = SQL.replace("$1", "to_tsvector('zh_en', question) || to_tsvector('zh_en', answer)")
            SQL = SQL.replace("$2", f"{ids}")
            await Tortoise.get_connection("pgsql").execute_query(SQL)


class PgDocumentsModel(DbModel):
    uuid = fields.UUIDField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="拥有者ID")
    file_id = fields.IntField(null=False, default=0, description="文件的ID")
    vector_model = fields.CharField(null=False, max_length=200, default="", description="向量模型")
    vector_alias = fields.CharField(null=False, max_length=200, default="", description="向量别名")
    page_no = fields.IntField(null=False, default=0, description="分页页码")
    page_nv = fields.CharField(null=False, max_length=100, default="", description="分页范围")
    chunk_index = fields.IntField(null=False, default=0, description="分段下标")
    chunk_texts = fields.TextField(default="", description="分段内容")
    metadata = fields.TextField(default="", description="元信息json")
    phrases = fields.TextField(description="中文分词")
    embedding = fields.TextField(description="向量内容")
    dimension = fields.IntField(null=False, default=0, description="向量维度")
    use_points = fields.IntField(null=False, default=0, description="消耗积分")
    use_tokens = fields.IntField(null=False, default=0, description="消耗积分")
    error = fields.TextField(null=False, default="", description="错误信息")
    status = fields.SmallIntField(null=False, default=0, description="训练状态: [0=等待中, 1=训练中, 2=成功, 3=失败]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "向量文档表"
        table = DbModel.table_prefix("documents")

    @classmethod
    async def update_ts_vectory(cls, uuids: List[str]):
        """ 更新分词信息 """
        if uuids:
            ids: str = ",".join(f"'{x}'" for x in uuids)
            SQL = cls.filter(uuid__in=[""]).update(phrases="").sql()
            SQL = SQL.replace("$1", "to_tsvector('zh_en', chunk_texts)")
            SQL = SQL.replace("$2", f"{ids}")
            await Tortoise.get_connection("pgsql").execute_query(SQL)


class PgAttachmentModel(DbModel):
    uuid = fields.UUIDField(pk=True, description="主键")
    dataset_id = fields.UUIDField(description="数据ID")
    archive_id = fields.IntField(null=False, default=0, description="文档ID")
    know_id = fields.IntField(null=False, default=0, description="知识库ID")
    scene = fields.CharField(default="", max_length=10, description="场景: [know]")
    type = fields.CharField(default="", max_length=10, description="类型: [image,video,file]")
    name = fields.CharField(default="", max_length=250, description="附件名称")
    path = fields.CharField(default="", max_length=250, description="附件路径")
    ext = fields.CharField(default="", max_length=10, description="附件扩展")
    size = fields.IntField(null=False, default=0, description="附件大小")
    sort = fields.IntField(null=False, default=0, description="排序编号")
    signed = fields.CharField(null=False, max_length=32, default="", description="唯一签名")
    vector_model = fields.CharField(null=False, max_length=200, default="", description="向量模型")
    vector_alias = fields.CharField(null=False, max_length=200, default="", description="向量别名")
    question = fields.TextField(default="", description="附件描述")
    phrases = fields.TextField(description="中文分词")
    embedding = fields.TextField(description="向量内容")
    dimension = fields.IntField(null=False, default=0, description="向量维度")
    use_points = fields.IntField(null=False, default=0, description="消耗积分")
    use_tokens = fields.IntField(null=False, default=0, description="消耗积分")
    error = fields.TextField(null=False, default="", description="错误信息")
    status = fields.SmallIntField(null=False, default=0, description="学习状态: [1=排队中, 2=学习中, 3=成功, 4=失败]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "附件索引表"
        table = DbModel.table_prefix("attachment")

    @classmethod
    async def update_ts_vectory(cls, dataset_id: str):
        """ 更新分词信息 """
        lists = await cls.filter(dataset_id=dataset_id).all().values("uuid")
        uuids = [item["uuid"] for item in lists]
        if uuids:
            ids: str = ",".join(f"'{x}'" for x in uuids)
            SQL = cls.filter(uuid__in=[""]).update(phrases="").sql()
            SQL = SQL.replace("$1", "to_tsvector('zh_en', question)")
            SQL = SQL.replace("$2", f"{ids}")
            await Tortoise.get_connection("pgsql").execute_query(SQL)
